///@file PolynomialChaos.cpp
///@brief Implementation of methods for polynomial chaos studies
///@author Arnaud Duval
///@version 0.0
///@date 2010/04/14

#include "PolynomialChaos.h"
#include <string>
#include <fstream>
#include <iostream>
#include <cstdlib>
#include <cmath>
#include "SobolSequence.h"
#include "FileUtils.h"
#include "StringUtils.h"
#include "../../mathobject/include/Matrix.h"
#include "PCbasis.h"

///@brief Preprocessor for polynomial chaos study
///@param deffile File containing the definition of the study
///@warning Must write a documentation for the structure of the definition file
void ChaosPrepro(char* deffile)
{
	std::ofstream fout("summary.txt");
	fout.precision(15);
	std::cout << "Bergamote design of experiment generator\n";
	int dim, nbpts;
	std::ifstream fin(deffile, std::ios::in);
	std::string buffer;
	std::getline(fin, buffer);
	fin >> dim;
	std::cout << "Number of parameters : " << dim << "\n";
	fout << "[Number of parameters]\n";
	fout << dim << "\n";
	std::getline(fin, buffer);
	std::getline(fin, buffer);
	fin >> nbpts;
	std::cout << "Number of experiences : " << nbpts << "\n";
	fout << "[Number of realisations]\n";
	fout << nbpts << "\n";
	std::getline(fin, buffer);
	std::getline(fin, buffer);
	SobolSequence seq(dim, nbpts);
	fout << "[Definition of parameters]\n";
	for(int i = 0 ; i < dim ; i++)
	{
		std::string type;
		fin >> type;
		if(type == "gauss")
		{
			double mu, sig;
			fin >> mu;
			fin >> sig;
			std::cout << "Parameter " << i+1 << " is a Gauss distribution\n";
			std::cout << "\t\tAverage value : " << mu << "\n";
			std::cout << "\t\tStandard deviation : " << sig << "\n";
			seq.TransformGauss(i+1, mu, sig);
			fout << type << "\t" << mu << "\t" << sig << "\n";
		}
		else if(type == "uniform")
		{
			double a, b;
			fin >> a;
			fin >> b;
			std::cout << "Parameter " << i+1 << " is a uniform distribution\n";
			std::cout << "\t\tLower bound : " << a << "\n";
			std::cout << "\t\tUpper bound : " << b << "\n";
			seq.TransformUniform(i+1, a, b);
			fout << type << "\t" << a << "\t" << b << "\n";
		}
		std::getline(fin, buffer);
	}
	std::getline(fin, buffer);
	unsigned int nbFiles;
	fin >> nbFiles;
	std::cout << "Number of input files to treat : " << nbFiles << "\n";
	std::getline(fin, buffer);
	std::getline(fin, buffer);
	std::string *fileNames = new std::string[nbFiles];
	std::string *fileExts = new std::string[nbFiles];
	for(unsigned int i = 0 ; i < nbFiles ; i++)
	{
		fin >> fileNames[i];
		fin >> fileExts[i];
		std::getline(fin, buffer);
		std::cout << "\t\t" << fileNames[i]+"."+fileExts[i] << "\n";
	}
	fin.close();

	fout << "[Design of experiment]\n";

	for(int i = 0 ; i < nbpts ; i++)
	{
		for(int j = 0 ; j < dim ; j++)
		{
			fout << seq[i+1][j] << "\t";
		}
		fout << "\n";
		for(unsigned int j = 0 ; j < nbFiles ; j++)
		{
			std::system(("cp "+fileNames[j]+"."+fileExts[j]+" "+fileNames[j]+"_"+IntToString(i+1,4)+"."+fileExts[j]).c_str());
			ReplaceStringInFile(fileNames[j]+"_"+IntToString(i+1,4)+"."+fileExts[j],"@@idx@@", "_"+IntToString(i+1,4));
			for(int k = 0 ; k < dim ; k++)
			{
				std::string paramName = "@@param"+IntToString(k+1,1)+"@@";
				ReplaceStringInFile(fileNames[j]+"_"+IntToString(i+1,4)+"."+fileExts[j],paramName, DoubleToString(seq[i+1][k],15));
			}
			std::cout << "Generated file : " << fileNames[j]+"_"+IntToString(i+1,4)+"."+fileExts[j] << "\n";
		}
	}
	fout.close();
	std::cout << "Output written to file summary.txt\n";

	delete fileNames;
	delete fileExts;
}

///@brief Solver for polynomial chaos study
///@param deffile File containing the definition of the study
///@warning Must write a documentation file for the structure of the definition file
void ChaosSolver(char* deffile)
{
	std::ifstream fin(deffile);
	std::string buffer, sumfile, resfile;
	std::getline(fin, buffer);
	std::getline(fin, sumfile);
	std::getline(fin, buffer);
	std::getline(fin, resfile);
	fin.close();

	std::ifstream sum(sumfile.c_str());
	unsigned int M;
	std::getline(sum, buffer);
	sum >> M;
	std::getline(sum, buffer);
	std::getline(sum, buffer);
	unsigned int N;
	sum >> N;
	std::getline(sum, buffer);
	std::getline(sum, buffer);
	std::string *types;
	types = new std::string[M];
	for(unsigned int i = 0 ; i < M ; i++)
	{
		double temp;
		sum >> types[i];
		sum >> temp;
		sum >> temp;
		std::getline(sum, buffer);
	}
	std::getline(sum, buffer);
	Matrix reals(N,M);
	for(unsigned int i = 1 ; i <= N ; i++)
	{
		for(unsigned int j = 1; j <= M ; j++)
			sum >> reals(i,j);
	}
	sum.close();

	std::ifstream res(resfile.c_str());
	Vector results(N);
	for(unsigned int i = 0 ; i < N ; i++)
		res >> results[i+1];
	res.close();

	Vector w(M);
	for(unsigned int i = 1 ; i <= M ; i++)
		w[i] = 1.;
	double q = 1.0;
	unsigned int p = 0;
	double eps1 = 1.E-5;
	double eps2 = 5.E-5;
	double critRfinal = 0.99;
	std::cout << "eps1 = " << eps1 << ", eps2 = " << eps2 << ", Rfinal = " << critRfinal << ", q = " << q << "\n";

	IndexSet A;		//Current Indexset
	IndexSet old;	//Sets already tested
	bool OK = false;
	Vector aa; //PC expansion coefficients (gestion of variables should be improved)
	while(!OK)
	{
		p++;
		std::cout << "p = " << p << ", ";
		IndexSet terms(M, p);
		terms.Truncate_qwNorm(p, q, w);
		terms = terms + A;
		terms = terms - old;
		//Traiter ici les termes deja elimines
		PCbasis basis(terms, types);
		Vector a = ComputePCCoefficients(reals, results, basis);
		Vector sensibility = ComputeACPSensibility(basis, a);
		w = ComputeACPWeights(sensibility);
		double R = DeterminationCoefficient(basis, a, reals, results);
		//std::cout << "R : " << R << "\n";
		double *R2var = new double[terms.Cardinality()];
		for(int i = 0 ; i < terms.Cardinality() ; i++)
		{
			IndexSet tempterms = terms;
			tempterms.Remove(i);
			PCbasis tempbasis(tempterms, types);
			Vector tempa = ComputePCCoefficients(reals, results, tempbasis);
			R2var[i] = DeterminationCoefficient(tempbasis, tempa, reals, results);
		}

		IndexSet toadd = terms;
		IndexSet toremove = terms;
		unsigned int addcpt = 0;
		unsigned int removecpy = 0;

		for(int i = terms.Cardinality() - 1 ; i >= 0 ; i--)
		{
			//Le terme augmente la determination
			if((R2var[i] > R)&&(R2var[i] - R > eps1))
			{
				toremove.Remove(i);
			}
			//Le terme fait baisser la determination
			else if((R2var[i] < R)&&(R - R2var[i] < eps2))
			{
				toadd.Remove(i);
			}
		}

		old = old + toremove;
		A = A + toadd;

		std::cout << "Card(A) = "<< A.Cardinality() << ", ";
		//Mesure finale de determination avec A
		PCbasis basisfinal(A, types);
		Vector afinal = ComputePCCoefficients(reals, results, basisfinal);
		double Rfinal = DeterminationCoefficient(basisfinal, afinal, reals, results);
		std::cout << "R = " << Rfinal << "\n";
		if (Rfinal > critRfinal)
			OK = true;

		delete R2var;
		aa = afinal;
	}
	//Write results into a file

	//Note : il faut ecrire a final (Vector), A (IndexSet) et les types. Ceci permet par la suite de reconstruire la base PC
	std::ofstream resPC("PCExpansion.pce");	//Must be improved to allow the user to choose file name	
	resPC << M << "\n";
	for(unsigned int i = 0 ; i < M ; i++)
	{
		resPC << types[i] << "\n";
	}
	resPC << aa.Dim() << "\n";
	resPC << aa << "\n";
	resPC << A.Cardinality() << "\n";
	resPC << A << "\n";
	delete[] types;
}

///@brief Compute Polynomial Chaos coefficients, giving a PC basis and a design of experiment
///@param reals design of experiment
///@param results results of experiment
///@param basis Polynomial Chaos basis
Vector ComputePCCoefficients(Matrix reals, Vector results, PCbasis basis)
{
	unsigned int N = reals.NbRows();
	Vector a(basis.Cardinality());
	Matrix Psi(N, basis.Cardinality());
	for(unsigned int i = 1 ; i <= N ; i++)
	{
		Psi[i] = basis.Evaluate(reals[i]);
	}
	a = (((Psi.Transpose()|Psi).Invert())|Psi.Transpose())|results;
	return a;
}

///@brief Compute total sensibilitis from PC expansion
///@param basis Polynomial Chaos basis
///@param a Coefficient of PC expansion
Vector ComputeACPSensibility(PCbasis basis, Vector a)
{
	unsigned int M = basis._M();
	Vector sensibility(M);
	for(unsigned int i = 1 ; i <= M ; i++)
	{
		sensibility[i] = 0;
		IndexSet I = basis.A().I_i(i);
		for(int j = 0 ; j < I.Cardinality() ; j++)
		{
			unsigned int k = 0;
			while(!(I[j] == (basis.A()[k])))
			{
				k++;
			}
			sensibility[i] += pow(a[k+1],2.);
		}
	}
	return sensibility;
}

///@brief Compute weights for adaptatif sparse anisotropic polynomial chaos
///@param sensibility Vector containing the total sensibilities 
Vector ComputeACPWeights(Vector sensibility)
{
	unsigned int M = sensibility.Dim();
	Vector Sk_Si(M);
	Vector w(M);
	for(unsigned int i = 1 ; i <= M ; i++)
	{
		for(unsigned int k = 1 ; k <= M ; k++)
			Sk_Si[k] = sensibility[k] - sensibility[i];
		w[i] = (sensibility.Sum() - Sk_Si.Max())/sensibility.Sum();
	}
	return w;
}

///@brief Compute determination coefficients of a PC expansion
///@param basis Polynomial chaos basis
///@param a Coefficients of PC expansion
///@param reals Design of experiment
///@param results experimental results
double DeterminationCoefficient(PCbasis basis, Vector a, Matrix reals, Vector results)
{
	unsigned int N = results.Dim();
	unsigned int M = basis._M();
	double R;
	double SCE = 0.;
	double SCR = 0.;
	double SCT = 0.;
	double ymoy = results.Average();

	Vector estimresults(N);
	for(unsigned int i = 1 ; i <= N ; i++)
		estimresults[i] = a|basis.Evaluate(reals[i]);

	for(unsigned int i = 1 ; i <= N ; i++)
	{
		SCE += pow(estimresults[i] - ymoy, 2.);
		SCR += pow(results[i] - estimresults[i], 2.);
		SCT += pow(results[i] - ymoy, 2.);
	}
	
	R =  1. - (SCR/(N - M - 1.))/(SCT/(N - 1.));
	return R;
}